import os, sqlite3, pandas as pd, requests
from sqlalchemy import create_engine, text
from dotenv import load_dotenv
load_dotenv()

class SQLRunner:
    def __init__(self, db_url: str):
        self.engine = create_engine(db_url)
    def exec(self, sql: str) -> pd.DataFrame:
        with self.engine.connect() as conn:
            return pd.read_sql_query(text(sql), conn)
    def init_from_csv(self, csv_path: str, table_name: str = "sales"):
        df = pd.read_csv(csv_path)
        with self.engine.connect() as conn:
            df.to_sql(table_name, conn, if_exists="replace", index=False)

class NL2SQLAgent:
    def __init__(self, schema_hint: str = "", provider=None, host=None, model=None,api_key=None):
        self.schema_hint = schema_hint
        self.provider = provider or os.getenv("LLM_PROVIDER", "ollama")
        self.host = host or os.getenv("LLM_HOST", "http://localhost:11434")
        self.model = model or os.getenv("LLM_MODEL", "deepseek-r1:1.5b")
        self.api_key = api_key or os.getenv("LLM_API_KEY", "")

    def _ask(self, prompt: str) -> str:
        if self.provider == "ollama":
            r = requests.post(f"{self.host}/api/generate", json={"model": self.model, "prompt": prompt, "stream": False})
            r.raise_for_status()
            return r.json().get("response", "")
        else:
            # 以 DashScope 为例，其它厂商同理
            headers = {"Authorization": f"Bearer {self.api_key}","Content-Type": "application/json"}
            r = requests.post(f"{self.host}/chat/completions", headers=headers, json={"model": self.model, "messages":[{"role":"user","content":prompt}]})
            r.raise_for_status()
            return r.json()["choices"][0]["message"]["content"]

    def generate_sql(self, question: str) -> str:
        prompt = f"""你是SQL生成助手。根据以下数据库结构和问题，生成可在对应数据库上直接执行的SQL。

        只返回SQL，不要任何解释或注释。
        [数据库结构]{self.schema_hint}
        [问题]{question}
        """

        sql = self._ask(prompt).strip().strip("```").replace("sql", "")
        return sql